"""
https://github.com/garythung/trashnet/blob/master/data/dataset-resized.zip
"""
import numpy as np
from keras.preprocessing.image import ImageDataGenerator, load_img, img_to_array
from keras.applications.efficientnet_v2 import EfficientNetV2B0, preprocess_input, decode_predictions
from keras.layers import Dense, Flatten, BatchNormalization
from keras.models import Sequential, Model, load_model
from keras.callbacks import EarlyStopping

img_size = (64, 64)
train_dir = '/Users/summy/Downloads/dataset-resized'


def build_model(pretrained_model, num_classes):
	return Sequential([
		pretrained_model,
		Flatten(),
		# BatchNormalization(),
		Dense(1024, activation='relu'),
		BatchNormalization(),
		Dense(num_classes, activation='softmax')
	])


def build_model2(pretrained_model, num_classes):
	head = Sequential([
		Flatten(),
		# BatchNormalization(),
		Dense(1024, activation='relu'),
		BatchNormalization(),
		Dense(num_classes, activation='softmax')
	])
	output = head(pretrained_model.output)
	return Model(inputs=pretrained_model.input, outputs=output)


def predict(X_pred):
	model = load_model('model.keras')
	preds = model.predict(X_pred)
	print(preds)
	

def predict2(pretrained_model, X_pred):
	model = build_model2(pretrained_model, 6)
	preds = model.predict(X_pred)
	print(preds)
	model.load_weights('model_weights.keras')
	preds = model.predict(X_pred)
	print(preds)
	

if __name__ == '__main__':
	pretrained = EfficientNetV2B0(include_top=False, weights='imagenet', input_shape=(*img_size, 3))
	for layer in pretrained.layers[:15]:
		layer.trainable = False

	datagen = ImageDataGenerator(
		preprocessing_function=preprocess_input,
		rotation_range=20,
		width_shift_range=0.2,
		height_shift_range=0.2,
		horizontal_flip=True,
		validation_split=0.2)

	dataset = datagen.flow_from_directory(
		train_dir,
		target_size=img_size,
		batch_size=16,
		class_mode='categorical')

	# print(pretrained.output_shape)  # (None, 2, 2, 1280)
	model = build_model(pretrained, dataset.num_classes)
	model.layers[0].trainable = False

	# model = build_model2(pretrained, dataset.num_classes)
	# print(model.summary())
	model.compile(loss='categorical_crossentropy', optimizer='adam', metrics=['accuracy'])
	early_stopping = EarlyStopping(monitor='loss', patience=3)
	model.fit(dataset, epochs=1, callbacks=early_stopping)
	# model.save('model.keras')
	model.save('model2')  # 如果转换为onnx，保存的是文件夹
	# model.save_weights('model_weights.keras')
	
	img = load_img('/Users/summy/Downloads/dataset-resized/glass/glass1.jpg', target_size=img_size)
	print(img.size)
	x = img_to_array(img)
	x = np.expand_dims(x, axis=0)
	x = preprocess_input(x)
	
	predict(x)
	# predict2(pretrained, x)

